// moon: The build system and package manager for MoonBit.
// Copyright (C) 2024 International Digital Economy Academy
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
// For inquiries, you can contact us via e-mail at jichuruanjian@idea.edu.cn.

//! Constructs the build plan from the given build conditions.
//!
//! This module handles a different thing from the package dependency graph
//! generated by [`crate::solve`]. `solve` mainly constructs a dependency graph
//! that is not affected by what we're actually building, nor does it cares
//! about what actions are actually carried out on the dependency graph (like
//! check, build or test).
//!
//! In contrast, this module generates the actual build plan from a list of
//! input action nodes. Irrelevant packages are not included in this graph,
//! nor does irrelevant actions.
//!
//! The result of this module is analogous to Rust Cargo's [Compile unit graph][cu].
//! The reason why this is two graphs here in MoonBuild and only one graph in
//! Cargo is because our build procedure is a little more verbose than Cargo.
//!
//! [cu]: https://github.com/rust-lang/cargo/blob/master/src/cargo/core/compiler/unit.rs

use std::{collections::HashMap, path::PathBuf};

use indexmap::{set::MutableValues, IndexSet};
use log::{debug, info, trace};
use moonutil::{
    common::TargetBackend,
    cond_expr::{OptLevel, ParseCondExprError},
};
use petgraph::{prelude::DiGraphMap, visit::DfsPostOrder};

use crate::{
    cond_comp::{self, CompileCondition},
    discover::DiscoverResult,
    model::{BuildPlanNode, BuildTarget, TargetAction, TargetKind},
    pkg_solve::DepRelationship,
};

/// A directed graph representation of build dependencies and targets.
///
/// `BuildPlan` maintains a directed graph where nodes represent build
/// components and edges represent dependencies between them. It also stores
/// specifications for each build target, mapping targets to their detailed
/// configuration and requirements.
#[derive(Default)]
pub struct BuildPlan {
    /// The build dependency graph.
    ///
    /// Each node in this graph represents a step in building, and edges
    /// represent the dependencies between build steps, pointing **from each
    /// step to what it depends on**.
    graph: DiGraphMap<BuildPlanNode, ()>,

    /// The specification of each build target.
    spec: HashMap<BuildPlanNode, BuildActionSpec>,

    /// The nodes that were used as input to the build plan.
    input_nodes: Vec<BuildPlanNode>,
}

impl BuildPlan {
    pub fn get_spec(&self, node: BuildPlanNode) -> Option<&BuildActionSpec> {
        self.spec.get(&node)
    }

    /// Get the list of nodes that **the given node depends on**.
    pub fn dependency_nodes(
        &self,
        node: BuildPlanNode,
    ) -> impl Iterator<Item = BuildPlanNode> + '_ {
        self.graph
            .neighbors_directed(node, petgraph::Direction::Outgoing)
    }

    /// Get the list of nodes that **depend on the given node**.
    pub fn consumer_nodes(&self, node: BuildPlanNode) -> impl Iterator<Item = BuildPlanNode> + '_ {
        self.graph
            .neighbors_directed(node, petgraph::Direction::Incoming)
    }

    pub fn all_nodes(&self) -> impl Iterator<Item = BuildPlanNode> + '_ {
        self.graph.nodes()
    }

    pub fn node_count(&self) -> usize {
        self.graph.node_count()
    }

    pub fn input_nodes(&self) -> &[BuildPlanNode] {
        &self.input_nodes
    }
}

/// Common information about a moonbit package being built
#[derive(Debug)]
pub struct BuildTargetInfo {
    pub files: Vec<PathBuf>,
    // we currently don't need this, as it's controlled by build-wise options
    // /// Whether compiling this target needs the standard library
    // pub std: bool,

    // we currently don't need this, as it's directly copied from mod.json.
    // pub is_main: bool,
}

#[derive(Debug)]
pub struct LinkCoreInfo {
    /// The targets in **initialization order**.
    pub linked_order: Vec<BuildTarget>,
    // we currently don't need this, as it's controlled by build-wise options
    // /// Whether linking this target needs the standard library
    // pub std: bool,
}

/// The specification of the specific build target, e.g. its files.
#[derive(Debug)]
pub enum BuildActionSpec {
    /// Check the given list of MoonBit source files.
    ///
    /// Outgoing edges of this node represents the direct dependencies.
    Check(BuildTargetInfo),

    /// Build the given list of MoonBit source files.
    ///
    /// Outgoing edges of this node represents the direct dependencies.
    BuildMbt(BuildTargetInfo),

    /// Build the given list of C source files.
    BuildC(Vec<PathBuf>),

    /// Link the core files from the given list of targets, **in order**.
    ///
    /// Outgoing edges of this node represents the build targets it depends on,
    /// but there's another copy of it, **ordered**, in the value's payloads,
    /// since the build command will need to use it.
    LinkCore(LinkCoreInfo),

    /// Make executable; link with the C artifacts of the given list of targets,
    /// if any.
    MakeExecutable { link_c_stubs: Vec<BuildTarget> },

    /// Generate the MBTI file for the given package.
    GenerateMbti,

    /// Bundle the packages specified by the outgoing edges.
    Bundle,

    /// Generates the test driver for the given target
    GenerateTestDriver(BuildTargetInfo),
}

/// Represents the environment in which the build is being performed.
pub struct BuildEnvironment {
    // FIXME: Target backend should go into the solver, not here
    pub target_backend: TargetBackend,
    pub opt_level: OptLevel,
    /// Whether compiling requires the standard library.
    ///
    /// TODO: Move this to per-package/module.
    pub std: bool,
    // Can have more, e.g. cross compile
}

/// Represents errors that may occur during build graph construction.
///
/// TODO: Will we even meet errors during build graph construction?
#[derive(Debug, thiserror::Error)]
pub enum BuildPlanConstructError {
    // TODO: This parsing should be moved earlier into the pipeline
    #[error("Error when parsing conditional compilation expression of {node:?}: {err}")]
    ParseCondExprError {
        node: BuildPlanNode,
        err: ParseCondExprError,
    },
}

/// Construct an abstract build graph from the given packages and input actions.
pub fn build_plan(
    packages: &DiscoverResult,
    build_deps: &DepRelationship,
    build_env: &BuildEnvironment,
    input: &[BuildPlanNode],
) -> Result<BuildPlan, BuildPlanConstructError> {
    info!("Constructing build plan with {} input nodes", input.len());
    debug!(
        "Build environment: backend={:?}, opt_level={:?}",
        build_env.target_backend, build_env.opt_level
    );

    let mut constructor = BuildPlanConstructor::new(packages, build_deps, build_env);
    constructor.build(input)?;
    let result = constructor.finish();

    info!(
        "Build plan construction completed with {} total nodes",
        result.graph.node_count()
    );
    Ok(result)
}

/// The struct responsible for holding the states and dependencies used during
/// the construction of a build plan.
struct BuildPlanConstructor<'a> {
    // Input environment
    packages: &'a DiscoverResult,
    build_deps: &'a DepRelationship,
    build_env: &'a BuildEnvironment,

    /// The resulting build plan
    res: BuildPlan,

    /// Currently pending nodes that need to be processed.
    pending: Vec<BuildPlanNode>,
}

impl<'a> BuildPlanConstructor<'a> {
    fn new(
        packages: &'a DiscoverResult,
        build_deps: &'a DepRelationship,
        build_env: &'a BuildEnvironment,
    ) -> Self {
        Self {
            packages,
            build_deps,
            build_env,
            res: BuildPlan::default(),
            pending: Vec::new(),
        }
    }

    fn finish(self) -> BuildPlan {
        self.res
    }

    fn build(&mut self, input: &[BuildPlanNode]) -> Result<(), BuildPlanConstructError> {
        assert!(
            self.pending.is_empty(),
            "Pending nodes should be empty before starting the build"
        );

        // Add the input node to the pending list
        for i in input {
            self.need_node(*i);
        }
        self.res.input_nodes.extend_from_slice(input);

        while let Some(node) = self.pending.pop() {
            // check if the node is already resolved
            if self.res.spec.contains_key(&node) {
                // Already resolved, skip
                continue;
            }

            self.build_action_dependencies(node)?;
        }
        Ok(())
    }

    /// Tell the build graph that we need to calculate the graph portion of a
    /// new node. To deduplicate pending nodes, this should be called before
    /// adding relevant edges to the graph (since the latter will also add the
    /// node into the graph).
    fn need_node(&mut self, node: BuildPlanNode) {
        if !self.res.spec.contains_key(&node) {
            self.pending.push(node);
            self.res.graph.add_node(node);
        }
    }

    /// Tell the build graph that we need to calculate the graph portion of a
    /// new node, and return that node for later usage. See [`Self::need_node`].
    fn need(&mut self, target: BuildTarget, action: TargetAction) -> BuildPlanNode {
        let node = BuildPlanNode { target, action };
        self.need_node(node);
        node
    }

    /// Tell the build graph that the given node has been resolved into a
    /// concrete action specification.
    fn resolved_node(&mut self, node_id: BuildPlanNode, spec: BuildActionSpec) {
        debug_assert!(
            !self.res.spec.contains_key(&node_id),
            "Node {:?} should not be resolved twice",
            node_id
        );
        debug_assert!(
            self.res.graph.contains_node(node_id),
            "Node {:?} should be in the graph before resolving",
            node_id
        );
        // Insert the node into the spec map
        self.res.spec.insert(node_id, spec);
    }

    fn add_edge(&mut self, start: BuildPlanNode, end: BuildPlanNode) {
        self.res.graph.add_edge(start, end, ());
    }

    /// Calculate the build action's dependencies and insert relevant edges to the
    /// build action graph.
    fn build_action_dependencies(
        &mut self,
        node: BuildPlanNode,
    ) -> Result<(), BuildPlanConstructError> {
        match node.action {
            TargetAction::Check => self.build_check(node),
            TargetAction::Build => self.build_build(node),
            TargetAction::BuildCStubs => self.build_build_c_stubs(node),
            TargetAction::LinkCore => {
                panic!(
                    "Link core should not appear in the wild without \
                    accompanied by MakeExecutable. Anytime it is met in the \
                    pending list, it should be already resolved."
                )
            }
            TargetAction::MakeExecutable => self.build_make_exec_link_core(node),
            TargetAction::GenerateTestInfo => self.build_gen_test_info(node),
        }
    }

    fn target_info_of(
        &self,
        node: BuildPlanNode,
    ) -> Result<BuildTargetInfo, BuildPlanConstructError> {
        // Resolve the source files
        let source_files = self.resolve_mbt_files_for_node(node)?;
        Ok(BuildTargetInfo {
            files: source_files,
            // is_main: pkg.raw.is_main,
            // std: self.build_env.std, // TODO: move to per-package
        })
    }

    fn build_check(&mut self, node: BuildPlanNode) -> Result<(), BuildPlanConstructError> {
        // Check depends on `.mi` of all dependencies, which practically
        // means the Check of all dependencies.
        for dep in self
            .build_deps
            .dep_graph
            .neighbors_directed(node.target, petgraph::Direction::Outgoing)
        {
            let dep_node = self.need(dep, TargetAction::Check);
            self.add_edge(node, dep_node);
        }

        self.resolved_node(node, BuildActionSpec::Check(self.target_info_of(node)?));

        Ok(())
    }

    fn build_build(&mut self, node: BuildPlanNode) -> Result<(), BuildPlanConstructError> {
        // Build depends on `.mi`` of all dependencies. Although Check can
        // also emit `.mi` files, since we're building, this action actually
        // means we need to build all dependencies.
        self.need_node(node);
        for dep in self
            .build_deps
            .dep_graph
            .neighbors_directed(node.target, petgraph::Direction::Outgoing)
        {
            let dep_node = self.need(dep, TargetAction::Build);
            self.add_edge(node, dep_node);
        }

        // If the given target is a test, we will also need to generate the test driver.
        if node.target.kind.is_test() {
            let gen_test_info = BuildPlanNode {
                action: TargetAction::GenerateTestInfo,
                ..node
            };
            self.need_node(gen_test_info);
            self.add_edge(node, gen_test_info);
        }

        self.resolved_node(node, BuildActionSpec::BuildMbt(self.target_info_of(node)?));

        Ok(())
    }

    fn build_gen_test_info(&mut self, node: BuildPlanNode) -> Result<(), BuildPlanConstructError> {
        self.need_node(node);
        let target_info = self.target_info_of(node)?; // FIXME: cache this
        self.resolved_node(node, BuildActionSpec::GenerateTestDriver(target_info));
        Ok(())
    }

    fn resolve_mbt_files_for_node(
        &self,
        node: BuildPlanNode,
    ) -> Result<Vec<PathBuf>, BuildPlanConstructError> {
        // FIXME: Should we resolve test drivers' paths, or should we leave it
        // in the lowering phase? The path to the test driver depends on the
        // artifact layout, so we might not be able to do that here, unless we
        // add some kind of `SpecialFile::TestDriver` or something.
        let pkg = self.packages.get_package(node.target.package);
        let source_files = cond_comp::filter_files(
            &pkg.raw,
            &pkg.root_path,
            pkg.source_files.iter().map(|x| x.as_path()),
            &CompileCondition {
                optlevel: self.build_env.opt_level,
                test_kind: node.target.kind.into(),
                backend: self.build_env.target_backend,
            },
        )
        .map_err(|err| BuildPlanConstructError::ParseCondExprError { node, err })?;
        Ok(source_files)
    }

    fn build_build_c_stubs(&mut self, node: BuildPlanNode) -> Result<(), BuildPlanConstructError> {
        // Depends on nothing, but anyway needs to be inserted into the graph.
        self.need_node(node);

        // Resolve the C stub files
        let pkg = self.packages.get_package(node.target.package);
        let c_source = pkg.c_stub_files.clone();
        self.resolved_node(node, BuildActionSpec::BuildC(c_source));

        Ok(())
    }

    /// Performs the construction of two actions in consecutive: Make Executable
    /// and Link Core.
    ///
    /// The two actions are always created together (Link Core is always a
    /// direct dependency of Make Executable, and there's no other actions that
    /// depends on Link Core), and both actions require traversing through the
    /// list of dependencies, so it's better to create both nodes at once,
    /// instead of in separate functions.
    fn build_make_exec_link_core(
        &mut self,
        make_exec_node: BuildPlanNode,
    ) -> Result<(), BuildPlanConstructError> {
        /*
            Link-core requires traversing all output of the current package's
            all transitive dependencies, and emitting them in DFS post-order.

            There are a couple of replacements needed to be done when the
            traversal completes:
            - Whitebox tests need to replace the normal package in the
                dependency graph (at the same position as the normal package).
                This is technically a circular dependency but anyway :)
            - Virtual package overrides need to replace their overridden
                packages in the dependency graph. This is done by not adding
                virtual packages at all when collecting the targets.
                TODO: virtual packages are not yet implemented here.
        */

        debug!(
            "Building MakeExecutable for target: {:?}",
            make_exec_node.target
        );
        debug!("Performing DFS post-order traversal to collect dependencies");
        // This DFS is shared by both LinkCore and MakeExecutable actions.
        let mut dfs = DfsPostOrder::new(&self.build_deps.dep_graph, make_exec_node.target);
        // This is the link core sources
        let mut link_core_deps = IndexSet::new();
        // This is the C stub sources
        let mut c_stub_deps = Vec::new();
        // DFS itself
        while let Some(next) = dfs.next(&self.build_deps.dep_graph) {
            if next.kind == TargetKind::WhiteboxTest {
                // Replace whitebox tests, if any
                let source_target = next.package.build_target(TargetKind::Source);
                if let Some(source_idx) = link_core_deps.get_index_of(&source_target) {
                    let source_mut = link_core_deps
                        .get_index_mut2(source_idx)
                        .expect("Source index is valid");
                    *source_mut = next;
                    continue;
                } else {
                    // No source target found, resort to regular path
                }
            }

            // Regular package
            link_core_deps.insert(next);
            // If there's any C stubs, add it (native only)
            let pkg = self.packages.get_package(next.package);
            trace!("DFS post iterated: {}", pkg.fqn);
            if self.build_env.target_backend.is_native() && !pkg.c_stub_files.is_empty() {
                c_stub_deps.push(next);
            }
        }

        let link_core_node = self.need(make_exec_node.target, TargetAction::LinkCore);

        // Add edges to all dependencies
        // Note that we have already replaced unnecessary dependencies
        for target in &link_core_deps {
            let dep_node = BuildPlanNode {
                target: *target,
                action: TargetAction::Build,
            };
            self.need_node(dep_node);
            self.add_edge(link_core_node, dep_node);
        }

        let targets = link_core_deps.into_iter().collect();
        let spec = BuildActionSpec::LinkCore(LinkCoreInfo {
            linked_order: targets,
            // std: self.build_env.std, // TODO: move to per-package
        });
        self.resolved_node(link_core_node, spec);

        // Add edge from make exec to link core
        self.add_edge(make_exec_node, link_core_node);

        // Add dependencies of make exec
        for target in &c_stub_deps {
            let dep_node = self.need(*target, TargetAction::BuildCStubs);
            self.add_edge(make_exec_node, dep_node);
        }
        self.resolved_node(
            make_exec_node,
            BuildActionSpec::MakeExecutable {
                link_c_stubs: c_stub_deps,
            },
        );

        Ok(())
    }
}
